package com.study.basic

import org.apache.spark.sql.SparkSession

/**
 * 分层抽样(Stratified sampling)
 * sampleByKey 每次都通过给定的概率以一种类似于掷硬币的方式来决定这个观察值是否被放入样本，
 * 因此一遍就可以过滤完所有数据，最后得到一个近似大小的样本，但往往不够准确。
 *
 * sampleByKeyExtra 会对全量数据做采样计算。对于每个类别，其都会产生 （fk⋅nk）个样本，
 * 其中fk是键为k的样本类别采样的比例；nk是键k所拥有的样本数。
 * sampleByKeyExtra 采样的结果会更准确，有99.99%的置信度，但耗费的计算资源也更多。
 *
 * @author stephen
 * @date 2019-08-27 11:12
 */
object StratifiedSamplingDemo {

  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .appName(this.getClass.getSimpleName)
      .master("local[*]")
      .getOrCreate()

    spark.sparkContext.setLogLevel("warn")

    val data = spark.sparkContext.makeRDD(Array(
      ("female", "Lily"),
      ("female", "Lucy"),
      ("female", "Emily"),
      ("female", "Kate"),
      ("female", "Alice"),
      ("male", "Tom"),
      ("male", "Roy"),
      ("male", "David"),
      ("male", "Frank"),
      ("male", "Jack")))
    // 设置采取60%的female和40%的male
    val fractions: Map[String, Double] = Map("female" -> 0.6, "male" -> 0.4)
    // 用sampleByKey进行抽样
    val approxSample = data.sampleByKey(withReplacement = false, fractions, 1)
    // 输出抽样结果
    println("sampleByKey result: ")
    approxSample.collect().foreach(println)
    // 本应该抽取3个female和2个male，但结果抽取了5个female和1个male，结果并不够准确，
    // 不过在样本数量足够大且要求一定效率的情况下，用sampleByKey进行抽样还是可行的。

    val exactSample = data.sampleByKeyExact(withReplacement = false, fractions, 1)
    println("sampleByKeyExact result: ")
    exactSample.collect().foreach(println)
    // 所得结果和预想一致，但当样本数量比较大时，可能会耗时较久。

    // 参数说明
    //      withReplacement	每次抽样是否有放回
    //      fractions	控制不同key的抽样率
    //      seed	随机数种子

    spark.stop()
  }
}
